use fxhash::FxHashMap;
use mango_core::math::{Vec2, Vec3, Vec4,  Transform3D};

use super::super::{
    assets::ShaderLib,
    utils::{pre_processor::PreProcessor, StringUtil},
    Program, PropertyValue,
};

use super::vertex_attributes::VertexAttributeSize;

#[derive(Debug)]
pub struct ShaderReflectionAttribute {
    pub name: String,
    pub group: u32,
    pub location: u32,
    pub ty: String,
    pub value_ty: String,
    pub value: u32,
    pub size: u32,
    pub format: wgpu::VertexFormat,
}

#[derive(Debug, Clone)]
pub struct ShaderReflectionVarInfo {
    pub group: u32,
    pub binding: u32,
    pub var_ty: String,
    pub var_name: String,
    pub data_ty: String,
    pub date_is_builtin_type: bool,

    pub data_fields: Vec<ShaderReflectionStructInfo>,
}

impl Default for ShaderReflectionVarInfo {
    fn default() -> Self {
        Self {
            group: 0,
            binding: 0,
            var_ty: "".to_owned(),
            var_name: "".to_owned(),
            data_ty: "".to_owned(),
            date_is_builtin_type: false,
            data_fields: vec![],
        }
    }
}

#[derive(Debug, Clone)]
pub struct ShaderReflectionStructInfo {
    pub name: String,
    pub ty: String,
}


impl ShaderReflectionStructInfo {
    pub(crate) fn size_of(&self) -> usize {
        if self.ty == "vec2<f32>" {2}
        else if self.ty == "vec3<f32>" {3}
        else if self.ty == "vec4<f32>" {4}
        else if self.ty == "f32" {1}
        else if self.ty == "mat3x3<f32>" {9}
        else if self.ty == "mat4x4<f32>" {16}
        else {0}
    }

    pub(crate) fn into_variant(&self) -> PropertyValue {
        if self.ty == "vec2<f32>" { PropertyValue::Vec2(Vec2::ZERO) }
        else if self.ty == "vec3<f32>" { PropertyValue::Vec3(Vec3::ZERO) }
        else if self.ty == "vec4<f32>" { PropertyValue::Vec4(Vec4::ZERO) }
        else if self.ty == "f32" { PropertyValue::Float(0.0) }
        // else if self.ty == "mat3x3<f32>" { PropertyValue::Transform2D(Transform2D::identity()) }
        else if self.ty == "mat4x4<f32>" { PropertyValue::Transform3D(Transform3D::default()) }
        else { PropertyValue::Float(0.0) }
    }
}

#[derive(Default, Debug)]
pub struct ShaderReflection {
    // _shaderReflectionMap
    pub attributes: Vec<ShaderReflectionAttribute>,
    pub vs_variables: Vec<ShaderReflectionVarInfo>,
    pub fs_variables: Vec<ShaderReflectionVarInfo>,
    pub cs_variables: Vec<ShaderReflectionVarInfo>,
    pub groups: Vec<Vec<ShaderReflectionVarInfo>>,
    pub variables: FxHashMap<String, ShaderReflectionVarInfo>,
    pub use_split: bool,
}

impl ShaderReflection {
    pub fn parser(wgsl: String) -> Self {
        let mut reflection = Self::default();
        if let Some(_) = wgsl.find("@vertex") {
            reflection.attributes = Self::parser_vertex_old(&wgsl);
            reflection.vs_variables = Self::parse_variables(&wgsl);
        } else if let Some(_) = wgsl.find("@fragment") {
            reflection.fs_variables = Self::parse_variables(&wgsl);
        } else if let Some(_) = wgsl.find("@compute") {
            reflection.cs_variables = Self::parse_variables(&wgsl);
        }

        reflection
    }

    pub fn parse2(wgsl: &str, shader_base: &mut Program) {
        if shader_base.shader_reflection.is_none() {
            shader_base.shader_reflection = Some(Self::default());  
        }

        if let Some(reflection) = &mut shader_base.shader_reflection {
            if let Some(_) = wgsl.find("@vertex") {
                reflection.attributes = Self::parser_vertex(&shader_base.vs_entry_point,&wgsl);
                reflection.vs_variables = Self::parse_variables(&wgsl);
            } else if let Some(_) = wgsl.find("@fragment") {
                reflection.fs_variables = Self::parse_variables(&wgsl);
            } else if let Some(_) = wgsl.find("@compute") {
                reflection.cs_variables = Self::parse_variables(&wgsl);
            }
        }
    }

    pub fn combine_shader_reflection_var_info(
        groups: &mut Vec<Vec<ShaderReflectionVarInfo>>,
        variables: &mut FxHashMap<String, ShaderReflectionVarInfo>,
        shader_variables: &[ShaderReflectionVarInfo],
    ) {

        
        for info in shader_variables {
            let target_group = info.group as usize;
            let target_binding = info.binding as usize;
            if groups.get(target_group).is_none() && target_group >= groups.len() {
                let diff_count = info.group as usize - groups.len();
                groups.reserve(diff_count + 1);
                for _ in 0..=diff_count {
                    groups.push(vec![]);
                }
            }

            let combine_info = info.clone();
            if let Some(group) = groups.get_mut(target_group) {
                if target_binding >= group.len() {
                    let diff_count = target_binding - group.len();
                    group.reserve(diff_count + 1);
                    for _ in 0..=diff_count {
                        group.push(Default::default());
                    }
                }
                group[target_binding] = combine_info.clone();
                variables
                    .insert(info.var_name.to_owned(), combine_info);
            }
        }
        
   
    }

    pub fn final_apply(shader_base: &mut Program) {
        if let Some(reflection) = &mut shader_base.shader_reflection {
            Self::combine_shader_reflection_var_info(
                &mut reflection.groups,
                &mut reflection.variables,
                &reflection.vs_variables,
            );

            Self::combine_shader_reflection_var_info(
                &mut reflection.groups,
                &mut reflection.variables,
                &reflection.fs_variables,
            );
        }
    }

    pub fn get_shader_reflection2(
        shader_lib: &mut ShaderLib,
        code: &str,
        shader_base: &mut Program,
    ) {
        if !shader_base.shader_variant.is_empty() {
            let pre_shader = PreProcessor::parse(shader_lib, code, &mut shader_base.define_value);
            Self::parse2(&pre_shader, shader_base);
        }
    }

    pub fn parse_variables(wgsl: &str) -> Vec<ShaderReflectionVarInfo> {
        let mut position = 0;
        let mut variables = vec![];

        while position < wgsl.len() {
            if let Some(n_left_index) = StringUtil::str_find_at(wgsl, "@group(", position) {
                if let Some(n_right_index) = StringUtil::str_find_at(wgsl, ";", n_left_index) {
                    position = n_right_index;
                    let item = &wgsl[n_left_index..n_right_index];
                    let group = Self::extract(item, "@group(", ")");
                    let binding = Self::extract(item, "@binding(", ")");

                    let var_name;
                    let mut var_type = "var".to_owned();

                    if let Some(_) = item.find("var<") {
                        var_name = Self::extract(item, ">", ":");
                        var_type = Self::extract(item, "var<", ">")
                            .replace(",", "-")
                            .replace(" ", "");
                    } else {
                        var_name = Self::extract(item, "var", ":");
                    }
                    let last_index = item.rfind(":").unwrap() + 1;
                    let data_ty = item[last_index..].trim();
                    let date_is_builtin_type = Self::is_built_in_types(data_ty);
                    let data_fields = if date_is_builtin_type {
                        vec![]
                    } else {
                        Self::parse_struct_fields(wgsl, data_ty)
                    };

                    variables.push(ShaderReflectionVarInfo {
                        group: group.parse::<u32>().unwrap(),
                        binding: binding.parse::<u32>().unwrap(),
                        var_ty: var_type,
                        var_name,
                        data_ty: data_ty.to_owned(),
                        date_is_builtin_type,

                        data_fields,
                    });
                }
            } else {
                break;
            }
        }

        return variables;
    }

    fn extract(str: &str, left: &str, right: &str) -> String {
        let index_l = StringUtil::str_find_at(str, left, 0).unwrap_or_default() + left.len();
        let index_r = StringUtil::str_find_at(str, right, index_l).unwrap_or_default();

        // let index_l = str.find(left).unwrap() + left.len();
        // let temp = &str[index_l..];
        // let index_r = temp.find(right).unwrap();
        str[index_l..index_r].trim().to_owned()
    }

    fn is_built_in_types(date_ty: &str) -> bool {
        match date_ty {
            "i32" => true,
            "u32" => true,
            "f32" => true,
            v => {
                if let Some(i) = v.find('<') {
                    let ty = &v[0..i];
                    match ty {
                        "vec2" => true,
                        "vec3" => true,
                        "vec4" => true,
                        "mat3" => true,
                        "mat4" => true,
                        "array" => {
                            let last = v.rfind('>').unwrap();
                            Self::is_built_in_types(&v[(i + 1)..last])
                        }
                        _ => false,
                    }
                } else {
                    false
                }
            }
        }
    }

    fn parse_struct_fields(wgsl: &str, struct_name: &str) -> Vec<ShaderReflectionStructInfo> {
        let mut result = vec![];
        let mut position = 0;

        while position < wgsl.len() {
            if let Some(n_left_index) = StringUtil::str_find_at(wgsl, "struct ", position) {
                if let Some(n_right_index) = StringUtil::str_find_at(wgsl, "{", n_left_index) {
                    position = n_right_index;
                    let name = wgsl[n_left_index + 6..n_right_index].trim();
                    if name == struct_name {
                        let n_left_index =
                            StringUtil::str_find_at(wgsl, "{", n_left_index).unwrap();
                        let n_right_index =
                            StringUtil::str_find_at(wgsl, "}", n_left_index).unwrap();
                        let items = &wgsl[(n_left_index + 1)..n_right_index];
                        let fields = items.split(',').collect::<Vec<_>>();
                        for field in fields {
                            if let Some(index) = field.find(':') {
                                result.push(ShaderReflectionStructInfo {
                                    name: field[0..index].trim().to_owned(),
                                    ty: field[index + 1..].trim().to_owned(),
                                })
                            }
                        }
                    }
                }
            } else {
                break;
            }
        }
        result
    }


    pub fn gen_render_shader_variant(effect:&mut Program) -> String {
        let mut variant = format!("RenderShader({},{})",&effect.vs_name,&effect.fs_name);
        variant.push('|');
        variant.push_str(&Self::gen_shader_variant(effect));
        variant.push('|');
        // variant.push_str(&format!("{:?}",effect.shader_state));
        variant
    }

    pub fn gen_shader_variant(effect:&mut Program) -> String {
        let mut variant = "".to_owned();
        for (key,_value) in effect.uniforms.0.iter() {
            variant.push_str(&format!("{}:",key));
        }
        variant.push('|');
        for (key,value) in effect.const_values.iter() {
            variant.push_str(&format!("{key}:{value}"));
        }

        variant.push('|');
        for (key,value) in effect.define_value.iter() {
            variant.push_str(&format!("{key}:{value}"));
        }

        variant
    
    }

    pub fn parser_vertex_old(wgsl: &str) -> Vec<ShaderReflectionAttribute> {
        let mut attributes = vec![];
        let list = wgsl.split("fn main(").collect::<Vec<_>>();

        let block = list[1].split("->").collect::<Vec<_>>();
        let block_list = block[0].split("@").collect::<Vec<_>>();
        if block_list.len() > 1 {
            for str in block {
                let space_char = regex::Regex::new(r"\s").unwrap();
                let code = space_char.replace_all(str, "");
                let code = code.replace(",", "");
                let code = code.replace("\n", "");
                let code: String = code.replace("  ", " ");
                println!("code = {}", code);
                Self::parse_attribute(&code, &mut attributes);
            }
        } else {
            let code1 = block[0].split(':').collect::<Vec<_>>();
            let code = code1[1];
            let code = code.replace("  ", "");
            let code = code.replace(" ", "");
            let code = code.replace(")", "");
            let code = wgsl.split(&format!("struct {}", code)).collect::<Vec<_>>()[1];

            let start = code.find('{');
            let end = code.find('}');
            if start.is_some() && end.is_some() {
                let start = start.unwrap();
                let end = end.unwrap();
                let code = &code[start..end];
                let block_list = code.split('@').collect::<Vec<_>>();
                for i in 1..block_list.len() {
                    let ele = block_list[i];
                    let space_char = regex::Regex::new(r"\s").unwrap();
                    let expr = space_char.replace_all(ele, " ");
                    let expr = expr.replace('\n', " ");
                    let expr = expr.split(',').collect::<Vec<_>>()[0];
                    let expr = expr.replace("  ", " ");
                    Self::parse_attribute(&expr, &mut attributes);
                }
            }
        }

        attributes
    }
    pub fn parser_vertex(entry_point: &str, wgsl: &str) -> Vec<ShaderReflectionAttribute> {
        let mut attributes = vec![];
        // println!("wgsl = {:?}",wgsl);
        // std::fs::write("wgsl_txt.wgsl", wgsl);
        let list = wgsl
            .split(&format!("fn {}(", entry_point))
            .collect::<Vec<_>>();


        let block = list[1].split("->").collect::<Vec<_>>();
        let block_list = block[0].split("@").collect::<Vec<_>>();
        if block_list.len() > 1 {
            for str in block {
                let space_char = regex::Regex::new(r"\s").unwrap();
                let code = space_char.replace_all(str, "");
                let code = code.replace(",", "");
                let code = code.replace("\n", "");
                let code: String = code.replace("  ", " ");
                // println!("code = {}", code);
                Self::parse_attribute(&code, &mut attributes);
            }
        } else {
            let code1 = block[0].split(':').collect::<Vec<_>>();
            let code = code1[1];
            let code = code.replace("  ", "");
            let code = code.replace(" ", "");
            let code = code.replace(")", "");
          
            let struct_name = code;
            let code = wgsl.split(&format!("struct {}", struct_name)).collect::<Vec<_>>();

            if code.len() == 1 {
                panic!("Struct ‘{}’ is not exist!",struct_name);
            }
            
            let code = code[1];

            let start = code.find('{');
            let end = code.find('}');
            if start.is_some() && end.is_some() {
                let start = start.unwrap();
                let end = end.unwrap();
                let code = &code[start..end];
                let block_list = code.split('@').collect::<Vec<_>>();
                for i in 1..block_list.len() {
                    let ele = block_list[i];
                    let space_char = regex::Regex::new(r"\s").unwrap();
                    let expr = space_char.replace_all(ele, " ");
                    let expr = expr.replace('\n', " ");
                    let expr = expr.split(',').collect::<Vec<_>>()[0];
                    let expr = expr.replace("  ", " ");
                    
                    Self::parse_attribute(&expr, &mut attributes);
                }
            }
        }

        attributes
    }

    pub fn parse_attribute(line: &str, attributes: &mut Vec<ShaderReflectionAttribute>) {
        let name;
        let group = 0;
        let location;
        let ty;
        let value_ty;
        let value = 0;
        let size;
        let format: wgpu::VertexFormat;

        let builtin = line.find("builtin");
        let local = line.find("location");
        let loca = regex::Regex::new(r"\((.+?)\)").unwrap();
        let caps = loca.captures(line).unwrap();

        location = caps
            .get(1)
            .unwrap()
            .as_str()
            .to_owned()
            .parse::<u32>()
            .unwrap_or(u32::MAX);

        let cc = line.split(":").collect::<Vec<_>>();
        let n = cc[0].split(" ").collect::<Vec<_>>()[1];
        let n = n.replace("  ", " ");
        let n = n.replace(" ", "");
        name = n;

        let val_ty = cc[1];
        let val_ty = val_ty.replace("  ", " ");
        let val_ty = val_ty.replace(" ", "");
        let val_ty = val_ty.replace("\r", "");
        let val_ty = val_ty.replace(")", "");
        value_ty = val_ty.replace(")", "");
        let attr_type = VertexAttributeSize::from(value_ty.as_str());
        size = attr_type.size();
        format = attr_type.into();

        if let Some(_) = builtin {
            ty = "builtin".to_owned();
        } else if let Some(_) = local {
            ty = "location".to_owned();
        } else {
            ty = "".to_owned()
        }

        attributes.push(ShaderReflectionAttribute {
            name,
            group,
            location,
            ty,
            value_ty,
            value,
            size,
            format,
        })
    }
}
